package com.fy.fyai.service.downloaddecorator.rate;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.io.InputStream;

/**
 * @author fwq
 * @description 令牌桶限流器实现（精确控制每秒传输速度）
 * @date 2025-11-03 21:05
 */
@Slf4j
@Component
public class TokenBucketRateLimiter implements RateLimiter {
    
    /**
     * 对输入流进行限流处理
     * 
     * @param inputStream 原始输入流
     * @param bytesPerSecond 每秒允许传输的字节数（必须大于0）
     * @return 限流后的输入流
     */
    @Override
    public InputStream limit(InputStream inputStream, long bytesPerSecond) throws IOException {
        if (!needLimit(bytesPerSecond)) {
            return inputStream;
        }
        
        return new RateLimitedInputStream(inputStream, bytesPerSecond);
    }
    
    /**
     * 限流输入流实现类（细粒度速率控制，平滑流量分布）
     * 使用基于实时速率的控制算法，每次读取少量数据并校验速率
     */
    private static class RateLimitedInputStream extends InputStream {
        private final InputStream delegate;
        private final long bytesPerSecond;
        
        // 速率控制相关
        private long totalBytesRead;           // 总读取字节数
        private long startTimeNanos;            // 开始时间
        private static final long NANOS_PER_SECOND = 1_000_000_000L;
        
        // 每次读取的块大小（细粒度控制，8KB）
        private static final int CHUNK_SIZE = 8 * 1024; // 8KB
        
        public RateLimitedInputStream(InputStream delegate, long bytesPerSecond) {
            this.delegate = delegate;
            this.bytesPerSecond = bytesPerSecond;
            this.totalBytesRead = 0;
            this.startTimeNanos = System.nanoTime();
        }
        
        @Override
        public int read() throws IOException {
            byte[] buffer = new byte[1];
            int result = read(buffer, 0, 1);
            return result == 1 ? buffer[0] & 0xFF : -1;
        }
        
        @Override
        public int read(byte[] b, int off, int len) throws IOException {
            if (len <= 0) {
                return 0;
            }
            
            // 细粒度控制：每次读取小块数据，实时校验速率
            int totalRead = 0;
            int remaining = len;
            int offset = off;
            
            while (remaining > 0) {
                // 每次读取的块大小（不超过CHUNK_SIZE）
                int chunkSize = Math.min(remaining, CHUNK_SIZE);
                
                // 检查当前速率，如果需要限流则等待
                checkAndWaitIfNeeded();
                
                // 读取数据块
                int bytesRead = delegate.read(b, offset, chunkSize);
                if (bytesRead <= 0) {
                    break; // EOF
                }
                
                // 更新统计
                totalBytesRead += bytesRead;
                totalRead += bytesRead;
                remaining -= bytesRead;
                offset += bytesRead;
            }
            
            return totalRead > 0 ? totalRead : -1;
        }
        
        /**
         * 检查当前速率并等待（细粒度控制）
         * 基于已读取的总字节数和经过的时间，计算当前平均速率
         */
        private void checkAndWaitIfNeeded() {
            long currentTimeNanos = System.nanoTime();
            long elapsedNanos = currentTimeNanos - startTimeNanos;
            
            // 如果时间太短，不进行限流检查（避免频繁计算）
            if (elapsedNanos < 10_000_000) { // 小于10毫秒
                return;
            }
            
            // 计算当前平均速率（字节/秒）
            double elapsedSeconds = (double) elapsedNanos / NANOS_PER_SECOND;
            double currentRate = totalBytesRead / elapsedSeconds;
            
            // 如果当前速率超过限制，需要等待
            if (currentRate > bytesPerSecond) {
                // 计算应该等待多久才能达到目标速率
                // 目标：totalBytesRead / (elapsedSeconds + waitSeconds) = bytesPerSecond
                // 即：totalBytesRead = bytesPerSecond * (elapsedSeconds + waitSeconds)
                // waitSeconds = (totalBytesRead / bytesPerSecond) - elapsedSeconds
                double waitSeconds = (totalBytesRead / (double) bytesPerSecond) - elapsedSeconds;
                
                if (waitSeconds > 0) {
                    // 转换为纳秒并等待
                    long waitNanos = (long) (waitSeconds * NANOS_PER_SECOND);
                    if (waitNanos > 0) {
                        try {
                            // 等待时间太长时，限制最大等待时间（避免长时间阻塞）
                            long maxWaitNanos = 100_000_000; // 最多等待100ms
                            waitNanos = Math.min(waitNanos, maxWaitNanos);
                            
                            long waitMillis = waitNanos / 1_000_000;
                            int waitNanosRemainder = (int) (waitNanos % 1_000_000);
                            Thread.sleep(waitMillis, waitNanosRemainder);
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                        }
                    }
                }
            }
            
            // 如果超过1秒，重置统计（避免累积误差）
            if (elapsedNanos >= NANOS_PER_SECOND) {
                // 保留部分数据，平滑过渡
                long excessNanos = elapsedNanos - NANOS_PER_SECOND;
                double excessRatio = (double) excessNanos / NANOS_PER_SECOND;
                
                // 保留超出部分的字节数
                totalBytesRead = (long) (totalBytesRead * excessRatio);
                startTimeNanos = currentTimeNanos - (long) (excessNanos);
            }
        }
        
        @Override
        public void close() throws IOException {
            delegate.close();
        }
        
        @Override
        public int available() throws IOException {
            return delegate.available();
        }
    }
}
